#include "pxt.h"

#include "esp_system.h"
#include "esp_partition.h"
// #include "spi_flash_mmap.h"
#include "esp_log.h"
#include "esp_err.h"

namespace pxt {

VMImage *vmImg;
TaskHandle_t userCodeTask;

static const char *TAG = "VMLOAD";

// for testing
#if 0
static uint8_t * loadAPP(){
    const size_t size = 256 * 1024; // 256 KB

    const esp_partition_t *partition = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_ANY, "ffat");
    if (partition == NULL) {
        ESP_LOGE(TAG, "Failed to find 'ffat' partition");
        return 0;
    }

    if (size > partition->size) {
        ESP_LOGE(TAG, "'ffat' partition size is smaller than required size");
        return 0;
    }

    const void *mapped_ptr;
    esp_partition_mmap_handle_t mmap_handle;

    esp_err_t ret = esp_partition_mmap(partition, 0, size, ESP_PARTITION_MMAP_DATA, &mapped_ptr, &mmap_handle);
    if (ret != ESP_OK) {
        ESP_LOGE(TAG, "Failed to map 'ffat' partition: %s", esp_err_to_name(ret));
        return 0;
    }

    ESP_LOGI(TAG, "Successfully mapped 256 KB from 'ffat' partition");

    const uint8_t *data = (const uint8_t *)mapped_ptr;

    uint8_t * app = (uint8_t *)heap_caps_malloc(size + 64, MALLOC_CAP_SPIRAM);
    if (app == NULL) {
        ESP_LOGE(TAG, "Failed to allocate memory for app");
        return 0;
    }

    printf("app: %p\n", app);
    // align to 16 bytes
    uint8_t * app_aligned = (uint8_t *)(((uintptr_t)app + 15) & ~15);

    memcpy(app_aligned, data, size);

    for (int i = 0; i < 16; i++) {
        ESP_LOGI(TAG, "Data[%d] = 0x%02X", i, app_aligned[i]);
    }

    esp_partition_munmap(mmap_handle);

    return (uint8_t *)app_aligned;
}
#endif


static void vmStartCore(uint8_t *data, unsigned len) {
    unloadVMImage(vmImg);
    vmImg = NULL;

    userCodeTask = xTaskGetCurrentTaskHandle();

    gcPreStartup();


    auto img = loadVMImage(data, len);
    if (img->errorCode) {
        dmesg("validation error %d at 0x%x", img->errorCode, img->errorOffset);
        return;
    } else {
        dmesg("Validation OK");
    }
    vmImg = img;

    gcStartup();

    globals = (TValue *)app_alloc(sizeof(TValue) * getNumGlobals());
    memset(globals, 0, sizeof(TValue) * getNumGlobals());

    initRuntime();
}

extern "C" void usb_init();

void vmStart() {
    memInfo();
    install_gpio0_handler();

    // start usb before loading the VM
    usb_init();
    VMImageSection * sect = NULL;
    VMImageHeader * hd = NULL;
    uint8_t *data = NULL;
    
    button_init();
    int ret = usb_pre_mount();

    // try boot.pxt64, then binary.pxt64, then built-in
    char * program = load_default_program();
    if (program){
        data = usb_load_pxt(program);
        free(program);
    }
    if (data == NULL){
        data = usb_load_pxt(NULL);
    }
    if (data == NULL){
        data = usb_load_pxt("binary.pxt64");
    }
    target_init();

    if (data){
        sect = (VMImageSection *)data;
        ESP_LOGI(TAG, "Loaded image from DISK %p", sect);
    }

    if (!sect) {
        ESP_LOGE(TAG, "No image found");
        sect = (VMImageSection *)PXT_EXPORTData[4];
        panicCode = 959;
    }

    

    // if (sect->type != SectionType::InfoHeader || hd->magic0 != VM_MAGIC0 || hd->imageSize < 256) {
    //     dmesg("invalid image at %p", sect);
    //     return;
    // }
    // TODO: load image from user selected program
    while (1){
        if (panicCode == 0 && sect){
            hd = (VMImageHeader *)sect->data;
            if (sect->type != SectionType::InfoHeader || hd->magic0 != VM_MAGIC0 || hd->imageSize < 256) {
                dmesg("invalid image at %p", sect);
            } else {
                vmStartCore((uint8_t *)sect, hd->imageSize);
            }
        }

        if (panicCode) {
            ESP_LOGI(TAG, "Panic %d", panicCode);
            if (panicCode > 950 && panicCode < 960){
                if (panicCode == 951){
                    // restart current app
                    panicCode = 0;
                }
                // usb writing new image
                // TODO: update status from tft
                while (panicCode == 952){
                    vTaskDelay(pdMS_TO_TICKS(200));
                    // screen_panic("USB Writing", 0);
                    int progress = usb_get_write_progress();
                    char buf[64];
                    snprintf(buf, 32, "USB Writing %d%%", progress);
                    screen_panic(buf, 0);
                }
                if (panicCode == 953){
                    vTaskDelay(pdMS_TO_TICKS(500));
                    // new image loaded
                    screen_panic("USB Loaded", 0);
                    sect = (VMImageSection *)usb_get_new_image(); // 1st section is image info
                    hd = (VMImageHeader *)sect->data; // vm tables start after info section
                    panicCode = 0;
                }
                if (panicCode == 959){
                    screen_panic("USB Disk", 0);
                }
            } else {
                char buf[64];
                snprintf(buf, 32, "Panic %d", panicCode);
                screen_panic(buf, 0x1800);
            }
            vTaskDelay(pdMS_TO_TICKS(500));
        }
    }
}

} // namespace pxt